from enum import IntEnum

from lz4.frame import decompress as lz4_decompress, compress as lz4_compress
from typing_extensions import Buffer
from zstd import compress as zstd_compress, decompress as zstd_decompress


class CompressMethod(IntEnum):
    Plaintext = 0
    LZ4 = 1
    ZSTD = 2


def compress_blob(compress_method: CompressMethod, data: bytes) -> bytes:
    if compress_method == CompressMethod.Plaintext:
        return data
    elif compress_method == CompressMethod.LZ4:
        return compress_blob_lz4(data)
    elif compress_method == CompressMethod.ZSTD:
        return compress_blob_zstd(data)
    else:
        raise ValueError(f'unknown compression method {compress_method}')


def decompress_blob(compress_method: CompressMethod, data: Buffer) -> bytes:
    if compress_method == CompressMethod.Plaintext:
        return bytes(data)
    elif compress_method == CompressMethod.LZ4:
        return decompress_blob_lz4(data)
    elif compress_method == CompressMethod.ZSTD:
        return decompress_blob_zstd(data)
    else:
        raise ValueError(f'unknown compression method {compress_method}')


def compress_blob_lz4(data: Buffer) -> bytes:
    return lz4_compress(data)


def decompress_blob_lz4(data: Buffer) -> bytes:
    return lz4_decompress(data)


def compress_blob_zstd(data: Buffer) -> bytes:
    return zstd_compress(data, 1)


def decompress_blob_zstd(data: Buffer) -> bytes:
    return zstd_decompress(data)
